import socket
import logging

DEFAULT_IP = "127.0.0.1"
DEFAULT_PORT = 12345

log = logging

class  __BaseTcpSocket:
    def __init__(self, sock):
        self.__ip = DEFAULT_IP
        self.__port = DEFAULT_PORT
        self.__sock = sock

    def __del__(self):
        self.close()


    def remote(self, ip=None, port=None):
        if ip is not None and len(ip) > 1:
            self.__ip = ip
        if port is not None:
            self.__port = port
        return self

    def ip(self):
        return self.__ip
    def set_ip(self, val):
        if self.__ip != val:
            self.__ip = val
        return self
    def port(self):
        return self.__port
    def set_port(self, val):
        if self.__port != val:
            self.__port = val
        return self

    def connect(self):
        return self.__sock.connect((self.__ip, self.__port))

    def write(self, bytes_buffer):
        return self.__sock.send(bytes_buffer)

    def read(self, buf_size):
        return self.__sock.recv(buf_size)

    def disconnect(self):
        return self.__sock.shutdown(socket.SHUT_RDWR)

    def close(self):
        if self.__sock is not None:
            log.debug("BaseTcpSocket:[%s] close", self)
            self.__sock.close()
            self.__sock = None

    def local(self):
        return self.__sock.getsockname()

    def handle(self):
        return self.__sock

    def __str__(self):
        return "[{}] -> [{}]".format(self.local(), (self.__ip, self.__port))


class BlockTcpSocket(__BaseTcpSocket):
    def __init__(self):
        super().__init__(socket.socket(socket.AF_INET,socket.SOCK_STREAM))

    def __str__(self):
        return "Block Tcp:" + super().__str__()

class NonBlockTcpSocket(__BaseTcpSocket):
    def __init__(self):
        super().__init__(socket.socket(socket.AF_INET,socket.SOCK_STREAM| socket.SOCK_NONBLOCK))

    def __str__(self):
        return "Nonblock Tcp:" + super().__str__()



import unittest 
class TestSocket(unittest.TestCase):
    def test_block(self):
        trans = BlockTcpSocket()
        trans.connect()
        
        trans.write(f"hello from {trans}".encode('utf-8'))

        trans.disconnect()
        trans.close()

        del trans
        trans = None

    def test_nonblock(self):
        import selectors
        selector = selectors.DefaultSelector()

        trans = NonBlockTcpSocket()

        # connect-noblock
        try:
            log.info("nonblocking connect")
            trans.connect()
        except BlockingIOError:
            log.info("got blockingIOError when connecting")
            key = selector.register(trans.handle(), selectors.EVENT_WRITE, trans) 

            events = selector.select(1)
            if len(events) == 0:
                log.error("connect failed")
                self.assertTrue(False)
            else:
                for key, mask in events:
                    log.info("connected")
                    self.assertTrue(True)
        except InterruptedError:
            log.error("break by signal")
            self.assertTrue(False)
        else:
            log.error("connect error")
            self.assertTrue(False)

        try:
            if selector.get_key(trans.handle()):
                selector.unregister(trans.handle())
        except KeyError:
            self.assertTrue(False)
        else:
            pass

        # read
        key = selector.register(trans.handle(), selectors.EVENT_READ, trans) 
        events = selector.select(1)
        if len(events) == 0:
            self.assertTrue(True) #timeout -> OK
        else:  # data exist in channel
            for key, _mask in events:
                read_bytes = trans.read(32)
                self.assertTrue(len(read_bytes) > 0)
                log.info("read some bytes, len:%d, data:%s", len(read_bytes), read_bytes.decode('utf-8'))
        selector.unregister(trans.handle())

        # write
        start = 0
        msg = "hello world".encode('utf-8')
        while start < len(msg):
            key = selector.register(trans.handle(), selectors.EVENT_WRITE, trans) 

            events = selector.select(1)
            self.assertTrue(len(events) > 0)
            for key, mask in events:
                writed_size = trans.write(msg[start:])
                self.assertTrue(writed_size > 0)
                start += writed_size
        self.assertEqual(len(msg), start)
        log.info("writed bytes, expect msg len:%d, writed_size:%d", len(msg), start)
        selector.unregister(trans.handle())

        trans.disconnect()
        trans.close()

        selector.close()

if __name__ == "__main__":
    #logging.basicConfig(level=logging.DEBUG)
    unittest.main(verbosity=2)
